open Ir
open RegAllocation

(* 在 IR 步骤中，已经实现了同名隐藏(通过 id 号), ASM 步骤只需要翻译即可 *)
(* ASM 步骤已经有 变量名 信息(Var x 的 x), 还有 oprand 信息(运算结果在 Reg t 里面) *)
(* 但是 IR 步骤的 operand 是假设无穷多个寄存器的 *)
(* 这里需要考虑做寄存器分配和放在栈空间暂存 *)
(* 函数实现时要考虑在 call 之前保存使用寄存器, 因为可能在之后被更改 *)
(* 没有考虑到 IR 的虚拟寄存器号超过 t6; 此时不应该直接使用 *)

(* 写史山；完全没有必要的类型 *)
type reg_or_imm = R of string | I of int

(* 优化 IR 也使用这一偏移量 *)
let fix_stack_size = 1648
let stack_offset = ref 0
let var_env = Hashtbl.create 1024 (* 用于完全栈分配版本 *)

(* 优化阶段不使用 *)
let reg_map : (string, string) Hashtbl.t =
  Hashtbl.create 256 (* 维护变量名和分配的寄存器号的映射 *)

let spilled_vars : (operand, int) Hashtbl.t =
  Hashtbl.create 256 (* 被暂时存放到内存里的变量名和内存偏移量 *)

(* 函数名 -> 它会用到的 caller-saved 寄存器（除了 a0） *)
let func_used_callers : (string, string list) Hashtbl.t = Hashtbl.create 16
let reg_pool = ref [ "t0"; "t1"; "t2"; "t3"; "t4"; "t5"; "t6" ] (* 可用寄存器池 *)

let valid_regs =
  [
    "zero";
    "ra";
    "sp";
    "gp";
    "tp";
    "t0";
    "t1";
    "t2";
    "t3";
    "t4";
    "t5";
    "t6";
    "s0";
    "s1";
    "s2";
    "s3";
    "s4";
    "s5";
    "s6";
    "s7";
    "s8";
    "s9";
    "s10";
    "s11";
    "a0";
    "a1";
    "a2";
    "a3";
    "a4";
    "a5";
    "a6";
    "a7";
  ]

(* 一般变量使用 spill_to_stack *)
(* 仍然也是返回相对于 s0 的负偏移 *)
let try_spill_to_stack (op : operand) =
  (* 已经 spill 到某个地方了 *)
  (* 这里不检查是否有 spill 标记，因为调用者保护和被调用保护也是用这个函数; spill 标记只在需要写回(store_if_spilled)的时候判断 *)
  try Hashtbl.find spilled_vars op
  with _ ->
    stack_offset := !stack_offset + 4;
    if !stack_offset >= 1644 then failwith "to large!\n";
    Hashtbl.add spilled_vars op (- !stack_offset - 4);
    - !stack_offset - 4

let fresh_temp () : string =
  match !reg_pool with
  | r :: rest ->
      reg_pool := rest;
      r
  | [] -> failwith "No free temporary registers available"

(* 使用 reg 寄存器; 从 reg_pool 里面移除 *)
(* 仅仅更新 reg_pool *)
let use_temp reg =
  let new_pool = List.filter (fun (x : string) -> x <> reg) !reg_pool in
  reg_pool := new_pool

let free_temp reg =
  let to_remove = ref [] in
  (* 移除所有存放在 reg 里的 var *)
  Hashtbl.iter
    (fun var r -> if r = reg then to_remove := var :: !to_remove)
    reg_map;
  List.iter (Hashtbl.remove reg_map) !to_remove;
  reg_pool := reg :: !reg_pool

(* 优化的 IR 不使用这一函数 *)
let get_stack_offset var =
  try Hashtbl.find var_env var
  with Not_found -> failwith ("Unknown variable: " ^ var)

(* 相对于 s0 的负偏移 s0 = sp + fix_stack_size *)
(* > 8 的函数参数的内存地址是手动放入的, 并且不会再使用; 不需要进入 var_env *)
(* > 8 的函数参数是相对于 sp 的正偏移; 并且从 0 开始增加 *)
(* 直接 get 出来的就是一个负数 *)
let alloc_stack_standard var =
  try get_stack_offset var
  with _ ->
    stack_offset := !stack_offset + 4;
    Hashtbl.add var_env var (- !stack_offset - 4);
    (* 从 -8 开始分配 *)
    -(!stack_offset + 4)

(* 优化的 IR 不使用这一函数 *)
(* always not used *)
let operand_to_str = function
  | Reg r | Var r -> Printf.sprintf "%d(sp)" (get_stack_offset r)
  | Imm i -> Printf.sprintf "%d" i

(* 优化的 IR 不使用这一函数 *)
let load_operand (reg : string) (op : operand) : string =
  match op with
  | Imm i -> Printf.sprintf "\tli %s, %d\n" reg i
  | Reg r | Var r ->
      Printf.sprintf "\tlw %s, %d(s0)\n" reg
        (get_stack_offset (*alloc_stack_standard*) r)

(* 优化的 IR 不使用这一函数 *)
let compile_inst (inst : ir_inst) (need_stack : bool) : string =
  match inst with
  | Binop (op, dst, lhs, rhs) ->
      let dst_off =
        alloc_stack_standard
          (match dst with Reg r | Var r -> r | _ -> failwith "Bad dst")
      in
      let lhs_code = load_operand "t1" lhs in
      let rhs_code = load_operand "t2" rhs in
      let op_code =
        match op with
        | "+" -> "\tadd t0, t1, t2\n"
        | "-" -> "\tsub t0, t1, t2\n"
        | "*" -> "\tmul t0, t1, t2\n"
        | "/" -> "\tdiv t0, t1, t2\n"
        | "%" -> "\trem t0, t1, t2\n"
        | "==" -> "\tsub t0, t1, t2\n\tseqz t0, t0\n"
        | "!=" -> "\tsub t0, t1, t2\n\tsnez t0, t0\n"
        | "<=" -> "\tsgt t0, t1, t2\n\txori t0, t0, 1\n"
        | ">=" -> "\tslt t0, t1, t2\n\txori t0, t0, 1\n"
        | "<" -> "\tslt t0, t1, t2\n"
        | ">" -> "\tsgt t0, t1, t2\n"
        | "&&" -> "\tand t0, t1, t2\n"
        | "||" -> "\tor t0, t1, t2\n"
        | _ -> failwith ("Unknown binop: " ^ op)
      in
      lhs_code ^ rhs_code ^ op_code ^ Printf.sprintf "\tsw t0, %d(s0)\n" dst_off
  | Unop (op, dst, src) ->
      let dst_off =
        alloc_stack_standard
          (match dst with Reg r | Var r -> r | _ -> failwith "Bad dst")
      in
      let load_src = load_operand "t1" src in
      let op_code =
        match op with
        | "-" -> "\tneg t0, t1\n"
        | "!" -> "\tseqz t0, t1\n"
        | "+" -> "\tmv t0, t1\n"
        | _ -> failwith ("Unknown unop: " ^ op)
      in
      load_src ^ op_code ^ Printf.sprintf "\tsw t0, %d(s0)\n" dst_off
  | Assign (dst, src) ->
      let dst_off =
        alloc_stack_standard
          (match dst with Reg r | Var r -> r | _ -> failwith "Bad dst")
      in
      let load_src = load_operand "t0" src in
      load_src ^ Printf.sprintf "\tsw t0, %d(s0)\n" dst_off
  (* Not used *)
  | Load (dst, src) ->
      let dst_off =
        alloc_stack_standard
          (match dst with Reg r | Var r -> r | _ -> failwith "Bad dst")
      in
      let src_code = load_operand "t1" src in
      src_code ^ "\tlw t0, 0(t1)\n" ^ Printf.sprintf "\tsw t0, %d(s0)\n" dst_off
  (* Not used *)
  | Store (dst, src) ->
      let dst_code = load_operand "t1" dst in
      let src_code = load_operand "t2" src in
      dst_code ^ src_code ^ "\tsw t2, 0(t1)\n"
  | Call (dst, fname, args) ->
      let dst_off =
        alloc_stack_standard
          (match dst with Reg r | Var r -> r | _ -> failwith "Bad dst")
      in
      let args_code =
        List.mapi
          (fun i arg ->
            if i < 8 then load_operand (Printf.sprintf "a%d" i) arg
            else
              let offset = 4 * (i - 8) in
              load_operand "t0" arg ^ Printf.sprintf "\tsw t0, %d(sp)\n" offset)
          args
        |> String.concat ""
      in
      args_code ^ Printf.sprintf "\tcall %s\n\tsw a0, %d(s0)\n" fname dst_off
      (* 尾调用消除 *)
  | TailCall (fname, args) ->
      (* 1. 将 args 按顺序写回函数形参对应的栈位或寄存器 *)
      let param_stores =
        List.mapi
          (fun i arg ->
            (* 先把 arg 装到 t0，然后写到对应的参数槽 *)
            let load = load_operand "t0" arg in
            if i < 8 then
              (* a0..a7 *)
              load ^ Printf.sprintf "\tmv a%d, t0\n" i
            else
              (* 栈上传第 (i-8) 个溢出参数 *)
              let offset = 4 * (i - 8) in
              load ^ Printf.sprintf "\tsw t0, %d(s0)\n" offset)
          args
        |> String.concat ""
      in
      (* 2. 无需保存返回地址，直接跳回函数入口 *)
      let entry_label = fname ^ "_entry" in
      (* 这个标签名要跟你 func_to_ir 中插入的 entry_lbl 保持一致 *)
      param_stores ^ Printf.sprintf "\tj %s\n" entry_label
  | Ret None ->
      if need_stack then
        Printf.sprintf
          "\tlw ra, %d(s0)\n\tlw s0, %d(sp)\n\taddi sp, sp, %d\n\tret\n"
          (get_stack_offset "ra") (fix_stack_size - 4) fix_stack_size
      else "\tret\n"
  | Ret (Some op) ->
      let load_code = load_operand "a0" op in
      if need_stack then
        let ra_offset = alloc_stack_standard "ra" in
        load_code
        ^ Printf.sprintf
            "\tlw ra, %d(s0)\n\tlw s0, %d(sp)\n\taddi sp, sp, %d\n\tret\n"
            ra_offset (fix_stack_size - 4) fix_stack_size
      else load_code ^ "\tret\n"
  | Goto label -> Printf.sprintf "\tj %s\n" label
  | IfGoto (cond, label) -> (
      match cond with
      | Imm 0 -> ""
      | Imm _ -> Printf.sprintf "\tj %s\n" label
      | _ ->
          let cond_code = load_operand "t0" cond in
          cond_code ^ Printf.sprintf "\tbne t0, x0, %s\n" label)
  | Label label -> Printf.sprintf "%s:\n" label

(* 判断 imm 是否能放进 I 型指令 *)
let fits_imm12 i = i >= -2048 && i <= 2047

(* 如果不行，就降级到先 li 再算 *)
let emit_li_and_op instr dst src imm =
  (* 先把常量装到 t6，再执行寄存器-寄存器指令 *)
  Printf.sprintf "\tli t6, %d\n\t%s %s, %s, t6\n" imm instr dst src

let compile_inst_with_liveness (reg_map : string OperandHashtbl.t)
    (inst : ir_inst) (_live_out : OperandSet.t) (callee_restore_code : string)
    (need_stack : bool) : string =
  let temp_regs = ref [ "t5"; "t6"; "t4" ] in
  let temp_reg_map = Hashtbl.create 3 in

  let alloc_temp_reg (v : operand) : string =
    match !temp_regs with
    | [] -> failwith "Ran out of temporary registers for spilling"
    | r :: rest ->
        temp_regs := rest;
        Hashtbl.add temp_reg_map v r;
        r
  in

  (* get_dst_reg 返回：(加载代码, 实际寄存器) *)
  let get_dst_reg (op : operand) : string * string =
    match op with
    | Imm _ -> failwith "Imm can't be dst"
    | _ -> (
        match OperandHashtbl.find_opt reg_map op with
        | Some reg when reg <> "__SPILL__" -> ("", reg)
        | _ ->
            let tmp_reg = alloc_temp_reg op in
            let var_name =
              match op with
              | Var v -> "Var " ^ v
              | Reg v -> "Reg " ^ v
              | _ -> "?"
            in
            (* 兼容 allocate_stack; 参数分配使用 allocate_stack *)
            let offset, debug_info =
              try
                let op_string =
                  match op with Var x -> x | Reg x -> x | _ -> failwith ""
                in
                (get_stack_offset op_string, "alloc_var")
              with _ -> (try_spill_to_stack op, "spill_var")
            in
            ( Printf.sprintf "\tlw %s, %d(s0) # reload %s -- %s\n" tmp_reg
                offset var_name debug_info,
              tmp_reg ))
  in

  (* 只使用 t6 寄存器 *)
  let get_arg_reg (op : operand) : string * string =
    match op with
    | Imm _ -> failwith "Imm can't be dst"
    | _ -> (
        match OperandHashtbl.find_opt reg_map op with
        | Some reg when reg <> "__SPILL__" -> ("", reg)
        | _ ->
            let tmp_reg = "t6" in
            let var_name =
              match op with
              | Var v -> "Var " ^ v
              | Reg v -> "Reg " ^ v
              | _ -> "?"
            in
            (* 兼容 allocate_stack; 参数分配使用 allocate_stack *)
            let offset =
              try
                let op_string = match op with Var x -> x | _ -> failwith "" in
                get_stack_offset op_string
              with _ -> try_spill_to_stack op
            in
            ( Printf.sprintf "\tlw %s, %d(s0) # reload %s\n" tmp_reg offset
                var_name,
              tmp_reg ))
  in

  (* 需要考虑操作数为 Imm 情况 *)
  let get_src_reg (op : operand) : string * string =
    match op with
    | Imm i ->
        let tmp_reg = alloc_temp_reg op in
        (Printf.sprintf "\tli %s, %d # tmp load imm \n" tmp_reg i, tmp_reg)
    | _ -> (
        match OperandHashtbl.find_opt reg_map op with
        | Some reg when reg <> "__SPILL__" -> ("", reg)
        | _ ->
            let tmp_reg = alloc_temp_reg op in
            let var_name =
              match op with Var v -> "Var" ^ v | Reg v -> "Reg" ^ v | _ -> "?"
            in
            (* 兼容 allocate_stack; 参数分配使用 allocate_stack *)
            let offset, debug_info =
              try
                let op_string =
                  match op with Var x -> x | Reg x -> x | _ -> failwith ""
                in
                (get_stack_offset op_string, "alloc_var")
              with _ -> (try_spill_to_stack op, "spill_var")
            in
            ( Printf.sprintf "\tlw %s, %d(s0) # reload %s -- %s\n" tmp_reg
                offset var_name debug_info,
              tmp_reg ))
  in

  (* helper：如果 dst 是 spill，要生成 store 指令 *)
  let store_dst_if_spilled dst dst_reg =
    let get_stack_offset_opt op =
      (* 只针对放在栈上的参数 *)
      let op_string = match op with Reg _ | Imm _ -> "" | Var v -> v in
      try Some (get_stack_offset op_string) with _ -> None
    in
    let var_str =
      match dst with
      | Var v -> "Var " ^ v
      | Reg v -> "Reg " ^ v
      | _ -> "Cannot spill imm!"
    in
    match OperandHashtbl.find_opt reg_map dst with
    | Some "__SPILL__" ->
        let offset = try_spill_to_stack dst in
        Printf.sprintf "\tsw %s, %d(s0) # save spilled %s\n" dst_reg offset
          var_str
    | _ -> (
        (* 如果它分配过栈空间（如参数）也要写回去 *)
        match get_stack_offset_opt dst with
        | Some offset ->
            Printf.sprintf "\tsw %s, %d(s0) # writeback param %s\n" dst_reg
              offset var_str
        | None -> "")
  in

  let restore_stack_and_ret =
    if need_stack then
      Printf.sprintf
        "\tlw ra, %d(s0)\n\tlw s0, %d(sp)\n\taddi sp, sp, %d\n\tret\n"
        (get_stack_offset "ra") (fix_stack_size - 4) fix_stack_size
    else "\tret\n"
  in

  match inst with
  | TailCall (fname, args) ->
      (* 1) 先为所有参数生成装载/赋值代码 *)
      let param_setup =
        List.mapi
          (fun i arg ->
            (* load_arg 遵循你已有的 get_arg_reg 操作： *)
            let load_code, reg = get_arg_reg arg in
            if i < 8 then
              (* 用 a0…a7 传参 *)
              load_code ^ Printf.sprintf "\tmv a%d, %s\n" i reg
            else
              (* 溢出参数写到栈上，偏移与普通 Call 保持一致 *)
              let offset = 4 * (i - 8) in
              load_code ^ Printf.sprintf "\tsw %s, %d(s0)\n" reg offset)
          args
        |> String.concat ""
      in
      (* 2) 直接无条件跳到函数入口标签 *)
      (*    假定你在 IR-to-IR 阶段把入口 label 定名为 fname ^ "_entry" *)
      let entry_lbl = fname ^ "_entry" in
      param_setup ^ Printf.sprintf "\tj %s\n" entry_lbl
  | Binop (op, dst, lhs, rhs) -> (
      match dst with
      | Imm _ -> failwith "Binop dst 不应为 Imm"
      | _ -> (
          let dst_load, dst_reg = get_dst_reg dst in

          match (lhs, rhs) with
          (* ---------------------- 强度缩减优化 ----------------------- *)
          | _, Imm 0 when op = "+" || op = "-" ->
              (* x + 0 = x, x - 0 = x, x || 0 = x *)
              let lhs_load, r1 = get_src_reg lhs in
              let store_code = store_dst_if_spilled dst r1 in
              lhs_load ^ dst_load
              ^ Printf.sprintf "\tmv %s, %s\n" dst_reg r1
              ^ store_code
          | Imm 0, _ when op = "+" ->
              (* 0 + x = x, 0 || x = x *)
              let rhs_load, r2 = get_src_reg rhs in
              let store_code = store_dst_if_spilled dst r2 in
              rhs_load ^ dst_load
              ^ Printf.sprintf "\tmv %s, %s\n" dst_reg r2
              ^ store_code
          | _, Imm 1 when op = "*" || op = "/" ->
              (* x * 1 = x, x / 1 = x *)
              let lhs_load, r1 = get_src_reg lhs in
              let store_code = store_dst_if_spilled dst r1 in
              lhs_load ^ dst_load
              ^ Printf.sprintf "\tmv %s, %s\n" dst_reg r1
              ^ store_code
          | Imm 1, _ when op = "*" ->
              (* 1 * x = x *)
              let rhs_load, r2 = get_src_reg rhs in
              let store_code = store_dst_if_spilled dst r2 in
              rhs_load ^ dst_load
              ^ Printf.sprintf "\tmv %s, %s\n" dst_reg r2
              ^ store_code
          | _, Imm 0 when op = "*" ->
              (* x * 0 = 0 *)
              let lhs_load, _ = get_src_reg lhs in
              let store_code = store_dst_if_spilled dst "zero" in
              lhs_load ^ dst_load
              ^ Printf.sprintf "\tmv %s, zero\n" dst_reg
              ^ store_code
          | Imm 0, _ when op = "*" ->
              (* 0 * x = 0 *)
              let rhs_load, _ = get_src_reg rhs in
              let store_code = store_dst_if_spilled dst "zero" in
              rhs_load ^ dst_load
              ^ Printf.sprintf "\tmv %s, zero\n" dst_reg
              ^ store_code
          | _, Imm n when op = "*" && n land (n - 1) = 0 ->
              (* x * 2^n => slli *)
              let shift = int_of_float (log (float_of_int n) /. log 2.) in
              let lhs_load, r1 = get_src_reg lhs in
              let store_code = store_dst_if_spilled dst dst_reg in
              lhs_load ^ dst_load
              ^ Printf.sprintf "\tslli %s, %s, %d\n" dst_reg r1 shift
              ^ store_code
          | _, Imm n when op = "/" && n land (n - 1) = 0 ->
              (* x / 2^n => srai  *)
              let shift = int_of_float (log (float_of_int n) /. log 2.) in
              let lhs_load, r1 = get_src_reg lhs in
              let store_code = store_dst_if_spilled dst dst_reg in
              lhs_load ^ dst_load
              ^ Printf.sprintf "\tsrai  %s, %s, %d\n" dst_reg r1 shift
              ^ store_code
          (* ------------------ 指令简化：立即数运算 ------------------- *)
          | _, Imm n when op = "+" && fits_imm12 n ->
              let lhs_load, r1 = get_src_reg lhs in
              let store_code = store_dst_if_spilled dst dst_reg in
              lhs_load ^ dst_load
              ^ Printf.sprintf "\taddi %s, %s, %d\n" dst_reg r1 n
              ^ store_code
          | _, Imm n when op = "-" && fits_imm12 n ->
              let lhs_load, r1 = get_src_reg lhs in
              let store_code = store_dst_if_spilled dst dst_reg in
              lhs_load ^ dst_load
              ^ Printf.sprintf "\taddi %s, %s, %d\n" dst_reg r1 (-n)
              ^ store_code
          | _, Imm n when op = "<<" && fits_imm12 n ->
              let lhs_load, r1 = get_src_reg lhs in
              let store_code = store_dst_if_spilled dst dst_reg in
              lhs_load ^ dst_load
              ^ Printf.sprintf "\tslli %s, %s, %d\n" dst_reg r1 n
              ^ store_code
          | _, Imm n when op = ">>" && fits_imm12 n ->
              let lhs_load, r1 = get_src_reg lhs in
              let store_code = store_dst_if_spilled dst dst_reg in
              lhs_load ^ dst_load
              ^ Printf.sprintf "\tsrai  %s, %s, %d\n" dst_reg r1 n
              ^ store_code
          (* ------------------------ 默认路径 ------------------------ *)
          | _ ->
              let lhs_load, r1 = get_src_reg lhs in
              let rhs_load, r2 = get_src_reg rhs in
              let inst_code =
                match op with
                | "+" -> Printf.sprintf "\tadd %s, %s, %s\n" dst_reg r1 r2
                | "-" -> Printf.sprintf "\tsub %s, %s, %s\n" dst_reg r1 r2
                | "*" -> Printf.sprintf "\tmul %s, %s, %s\n" dst_reg r1 r2
                | "/" -> Printf.sprintf "\tdiv %s, %s, %s\n" dst_reg r1 r2
                | "%" -> Printf.sprintf "\trem %s, %s, %s\n" dst_reg r1 r2
                | "==" ->
                    Printf.sprintf "\tsub %s, %s, %s\n\tseqz %s, %s\n" dst_reg
                      r1 r2 dst_reg dst_reg
                | "!=" ->
                    Printf.sprintf "\tsub %s, %s, %s\n\tsnez %s, %s\n" dst_reg
                      r1 r2 dst_reg dst_reg
                | "<" -> Printf.sprintf "\tslt %s, %s, %s\n" dst_reg r1 r2
                | ">" -> Printf.sprintf "\tsgt %s, %s, %s\n" dst_reg r1 r2
                | "<=" ->
                    Printf.sprintf "\tsgt %s, %s, %s\n\txori %s, %s, 1\n"
                      dst_reg r1 r2 dst_reg dst_reg
                | ">=" ->
                    Printf.sprintf "\tslt %s, %s, %s\n\txori %s, %s, 1\n"
                      dst_reg r1 r2 dst_reg dst_reg
                | "&&" -> Printf.sprintf "\tand %s, %s, %s\n" dst_reg r1 r2
                | "||" -> Printf.sprintf "\tor %s, %s, %s\n" dst_reg r1 r2
                | _ -> failwith ("Unknown binop: " ^ op)
              in
              let store_code = store_dst_if_spilled dst dst_reg in
              String.concat ""
                [ lhs_load; rhs_load; dst_load; inst_code; store_code ]))
  | Unop (op, dst, src) ->
      let dst_load, dst_reg = get_dst_reg dst in
      let src_load, src_reg = get_src_reg src in
      let body =
        match op with
        | "-" -> Printf.sprintf "\tsub %s, zero, %s\n" dst_reg src_reg
        | "!" -> Printf.sprintf "\tseqz %s, %s\n" dst_reg src_reg
        | "+" -> "" (* noop *)
        | _ -> failwith ("Unknown unop: " ^ op)
      in
      let store_code = store_dst_if_spilled dst dst_reg in
      dst_load ^ src_load ^ body ^ store_code
  | Assign (dst, src) -> (
      let dst_load, dst_reg = get_dst_reg dst in
      match src with
      | Imm i ->
          let body = Printf.sprintf "\tli %s, %d\n" dst_reg i in
          let store_code = store_dst_if_spilled dst dst_reg in
          dst_load ^ body ^ store_code
      | _ ->
          (* 确定是寄存器类型 *)
          let src_load, src_reg = get_dst_reg src in
          let body = Printf.sprintf "\tmv %s, %s\n" dst_reg src_reg in
          let store_code = store_dst_if_spilled dst dst_reg in
          dst_load ^ src_load ^ body ^ store_code)
  | Load (dst, addr) ->
      (* 不会被使用到的分支; 先乱写一波 *)
      let dst_load, dst_reg = get_dst_reg dst in
      let addr_load, addr_reg = get_dst_reg addr in
      let body = Printf.sprintf "\tlw %s, 0(%s)\n" dst_reg addr_reg in
      let store_code = store_dst_if_spilled dst dst_reg in
      dst_load ^ addr_load ^ body ^ store_code
  | Store (addr, value) ->
      let addr_load, addr_reg = get_dst_reg addr in
      let val_load, val_reg = get_dst_reg value in
      addr_load ^ val_load ^ Printf.sprintf "\tsw %s, 0(%s)\n" val_reg addr_reg
  | Call (dst, fname, args) ->
      (* 从 func_used_callers 获取目标函数实际使用的 caller-saved 寄存器 *)
      let used_by_callee =
        match Hashtbl.find_opt func_used_callers fname with
        | Some regs -> regs
        | None -> [] (* 还没处理到被调用的函数，可选择保守处理 *)
      in

      let to_save = ref [] in

      OperandSet.iter
        (fun op ->
          match OperandHashtbl.find_opt reg_map op with
          | Some reg when reg <> "__SPILL__" && List.mem reg used_by_callee ->
              let offset = try_spill_to_stack op in
              to_save := (op, reg, offset) :: !to_save
          | _ -> ())
        _live_out;

      let save_code =
        List.map
          (fun (op, reg, offset) ->
            let var_str =
              match op with
              | Var v -> v
              | Reg v -> v
              | Imm i -> failwith (Printf.sprintf "%d" i)
            in
            Printf.sprintf "\tsw %s, %d(s0)  # save caller-saved %s\n" reg
              offset var_str)
          !to_save
        |> String.concat ""
      in

      let restore_code =
        List.map
          (fun (op, reg, offset) ->
            let var_str =
              match op with
              | Var v -> v
              | Reg v -> v
              | Imm i -> failwith (Printf.sprintf "%d" i)
            in
            Printf.sprintf "\tlw %s, %d(s0)  # restore caller-saved %s\n" reg
              offset var_str)
          !to_save
        |> String.concat ""
      in

      let temp_reg_to_be_moved = alloc_temp_reg (Var "__temp_moved") in
      let move_from_ret_to_temp = Printf.sprintf "\tmv %s, a0 # move a0 to temp_reg\n" temp_reg_to_be_moved in
      let _, dst_reg = get_dst_reg dst in
      let move_from_temp_to_dst =
        (* dst_load *)
        Printf.sprintf "\tmv %s, %s  # move return value from temp_reg %s\n" dst_reg temp_reg_to_be_moved temp_reg_to_be_moved
        ^ store_dst_if_spilled dst dst_reg
      in

      (* 设置函数参数 *)
      let arg_regs = [| "a0"; "a1"; "a2"; "a3"; "a4"; "a5"; "a6"; "a7" |] in
      let args_code =
        List.mapi
          (fun i arg ->
            match arg with
            | Imm n when i < 8 -> Printf.sprintf "\tli %s, %d\n" arg_regs.(i) n
            | _ when i < 8 ->
                let load, reg = get_arg_reg arg in
                load ^ Printf.sprintf "\tmv %s, %s\n" arg_regs.(i) reg
            | Imm n ->
                let off = 4 * (i - 8) in
                Printf.sprintf "\tli t6, %d\n\tsw t6, %d(sp)\n" n off
            | _ ->
                let off = 4 * (i - 8) in
                let load, reg = get_arg_reg arg in
                load ^ Printf.sprintf "\tsw %s, %d(sp)\n" reg off)
          args
        |> String.concat ""
      in

      (* 最终指令序列拼接 *)
      save_code ^ args_code
      ^ Printf.sprintf "\tcall %s\n" fname
      ^ move_from_ret_to_temp
      ^ restore_code 
      ^ move_from_temp_to_dst
      (* —— 新增：尾调用消除 —— *)
  | IfGoto (cond, label) -> (
      match cond with
      | Imm 0 -> "" (* 不跳转，省略整条指令 *)
      | Imm _ -> Printf.sprintf "\tj %s\n" label (* 恒跳转 *)
      | _ ->
          let cond_load, cond_reg = get_dst_reg cond in
          cond_load ^ Printf.sprintf "\tbnez %s, %s\n" cond_reg label)
  | Goto label -> Printf.sprintf "\tj %s\n" label
  | Label label -> Printf.sprintf "%s:\n" label
  | Ret None -> callee_restore_code ^ restore_stack_and_ret
  | Ret (Some op) ->
      (match op with
      | Imm i -> Printf.sprintf "\tli a0, %d\n" i
      | _ ->
          let ret_load, ret_reg = get_dst_reg op in
          let mv_code =
            if ret_reg = "a0" then ""
            else Printf.sprintf "\tmv a0, %s\n" ret_reg
          in
          ret_load ^ mv_code)
      ^ callee_restore_code ^ restore_stack_and_ret

(* let compile_block (blk : ir_block) (_ : string OperandHashtbl.t)
    (need_stack : bool) : string =
  blk.insts
  |> List.mapi (fun idx inst ->
         try compile_inst inst need_stack
         with exn ->
           let msg = Printexc.to_string exn in
           let inst_str =
             try Print_ir.string_of_ir_inst inst
             with _ -> "<unknown instruction>"
           in
           Printf.eprintf
             "Error compiling instruction at index %d in block %s:\n\
             \  %s\n\
             \  Reason: %s\n\
              %!"
             idx blk.label inst_str msg;
           raise exn)
  |> String.concat "" *)

(* let live = ref blk.live_out in

  (* 倒序遍历指令，并处理活跃变量 *)
  List.iter
    (fun inst ->
      let def, use = Optimazation.def_use_inst inst in
      let live_out = !live in
      let inst_code = compile_inst_with_liveness reg_map inst live_out in
      code_acc := inst_code :: !code_acc;
      (* 更新当前指令结束后的 live *)
      live := StringSet.union (StringSet.diff !live def) use)
    (List.rev blk.insts);

  String.concat "" !code_acc *)

(* 新的 compile_block 函数 *)
let compile_block (blk : ir_block) (reg_map : string OperandHashtbl.t)
    (need_stack : bool) (callee_restore_code : string) : string =
  let code_acc = ref [] in
  let live = ref blk.live_out in

  (* 倒序遍历指令，并处理活跃变量 *)
  List.iter
    (fun inst ->
      let def, use = Optimazation.def_use_inst inst in
      let live_out = !live in
      let inst_code =
        compile_inst_with_liveness reg_map inst live_out callee_restore_code
          need_stack
      in
      code_acc := inst_code :: !code_acc;
      (* 更新当前指令结束后的 live *)
      live := OperandSet.union (OperandSet.diff !live def) use)
    (List.rev blk.insts);

  String.concat "" !code_acc

let compile_func (f : ir_func) : string =
  Hashtbl.clear var_env;
  stack_offset := 0;

  (* 参数入栈 *)
  let param_setup_tp =
    Printf.sprintf "\tsw s0, %d(sp)\n" (fix_stack_size - 4)
  in

  (* 参数入栈 *)
  let param_setup =
    List.mapi
      (fun i name ->
        (* 前 8 个参数直接通过上面寄存器的映射访问; 后续参数放在栈上 *)
        if i >= 8 then
          (* -8(s0) 起分配; -4(s0) 存放的原 s0 *)
          let off = alloc_stack_standard name in
          (* 相对于 s0 的正偏移 *)
          Printf.sprintf "\tlw t5, %d(s0)\n\tsw t5, %d(s0)\n"
            (* offset 为 call 语句将第 i 个参数压入的偏移 *)
            (4 * (i - 8))
            off
        else
          let off = alloc_stack_standard name in
          Printf.sprintf "\tsw a%d, %d(s0)\n" i off)
      f.args
    |> String.concat ""
  in

  (* ra 入栈 *)
  let param_setup =
    param_setup ^ Printf.sprintf "\tsw ra, %d(s0)\n" (alloc_stack_standard "ra")
  in

  let body_code =
    f.body |> List.map (fun inst -> compile_inst inst true) |> String.concat ""
  in

  (* 检查 body_code 是否以 ret 结束; 没有默认添加 "\taddi sp, sp, 1024\n\tret\n" 语句; 其实可以前移到 IR 阶段 *)
  let body_code =
    if not (String.ends_with ~suffix:"\tret\n" body_code) then
      body_code
      ^ Printf.sprintf
          "\tlw ra, %d(s0)\n\tlw s0, %d(sp)\n\taddi sp, sp, %d\n\tret\n"
          (get_stack_offset "ra") (fix_stack_size - 4) fix_stack_size
    else body_code
  in
  let func_label = f.name in
  let prologue =
    Printf.sprintf "%s:\n\taddi sp, sp, -%d\n" func_label fix_stack_size
    ^ param_setup_tp
    ^ Printf.sprintf "\taddi s0, sp, %d\n" fix_stack_size
  in
  prologue ^ param_setup ^ body_code

let need_stack_frame (f : ir_func_o) (reg_map : string OperandHashtbl.t)
    (callee_saved : string list) : bool =
  let calls_other_functions =
    List.exists
      (fun blk ->
        List.exists
          (function
            | Call _ -> true
            | TailCall _ -> true (* tailcall 也要 stack 保存参数等 *)
            | _ -> false)
          blk.insts)
      f.blocks
  in
  let has_spilled_vars =
    OperandHashtbl.fold
      (fun _ reg acc -> acc || reg = "__SPILL__")
      reg_map false
  in
  calls_other_functions || has_spilled_vars || not (List.is_empty callee_saved)

let compile_func_o (f : ir_func_o) (print_alloc : bool) : string =
  Hashtbl.clear var_env;
  Hashtbl.clear reg_map;
  Hashtbl.clear spilled_vars;
  stack_offset := 0;

  Optimazation.liveness_analysis f.blocks print_alloc;
  let intervals = build_intervals f in
  (* true 表示启动打印 分配结果 *)
  let reg_map = linear_scan_allocate intervals print_alloc in

  (* Printf.printf "function %s need_stack %b\n" f.name need_stack; *)
  let used_caller_saved =
    let caller_saved =
      [
        "t0";
        "t1";
        "t2";
        "t3";
        "t4";
        "t5";
        "t6";
        "a1";
        "a2";
        "a3";
        "a4";
        "a5";
        "a6";
        "a7";
      ]
    in
    OperandHashtbl.fold
      (fun _ reg acc ->
        if List.mem reg caller_saved && not (List.mem reg acc) then reg :: acc
        else acc)
      reg_map []
  in
  let arg_regs = [| "a0"; "a1"; "a2"; "a3"; "a4"; "a5"; "a6"; "a7" |] in
  let param_regs =
    List.mapi (fun i _ -> if i < 8 then arg_regs.(i) else "") f.args
    |> List.filter (fun r -> r <> "")
  in

  let used_caller_saved =
    List.fold_left
      (fun acc reg -> if List.mem reg acc then acc else reg :: acc)
      used_caller_saved param_regs
  in
  Hashtbl.replace func_used_callers f.name used_caller_saved;

  let used_callee_saved =
    let callee_saved =
      [
        "s0"; "s1"; "s2"; "s3"; "s4"; "s5"; "s6"; "s7"; "s8"; "s9"; "s10"; "s11";
      ]
    in
    OperandHashtbl.fold
      (fun _ reg acc ->
        if List.mem reg callee_saved && not (List.mem reg acc) then reg :: acc
        else acc)
      reg_map []
  in

  let need_stack = need_stack_frame f reg_map used_callee_saved in

  let callee_save_code =
    used_callee_saved
    |> List.map (fun reg ->
           let offset = try_spill_to_stack (Reg reg) in
           Printf.sprintf "\tsw %s, %d(s0)  # save callee-saved\n" reg offset)
    |> String.concat ""
  in

  let callee_restore_code =
    used_callee_saved
    |> List.map (fun reg ->
           let offset = try_spill_to_stack (Reg reg) in
           Printf.sprintf "\tlw %s, %d(s0)  # restore callee-saved\n" reg offset)
    |> String.concat ""
  in

  (* prologue & param setup *)
  let prologue, param_setup =
    if need_stack then
      let param_setup_tp =
        Printf.sprintf "\tsw s0, %d(sp)\n" (fix_stack_size - 4)
      in
      let param_setup =
        List.mapi
          (fun i name ->
            if i >= 8 then
              let off = try_spill_to_stack (Var name) in
              Printf.sprintf "\tlw t5, %d(s0)\n\tsw t5, %d(s0)\n"
                (4 * (i - 8))
                off
            else
              (* OperandHashtbl.add reg_map (Var name) (Printf.sprintf "a%d" i); *)
              (* 感觉这个适配上确实有问题，先修改为使用 a_i *)
              let off = try_spill_to_stack (Var name) in
              Printf.sprintf "\tsw a%d, %d(s0)\n" i off
              (*let off = alloc_stack_standard name in
              Printf.sprintf "\tsw a%d, %d(s0)\n" i off*))
          f.args
        |> String.concat ""
      in
      let param_setup =
        param_setup
        ^ Printf.sprintf "\tsw ra, %d(s0)\n" (alloc_stack_standard "ra")
        ^ callee_save_code
      in
      let prologue =
        Printf.sprintf "%s:\n\taddi sp, sp, -%d\n" f.name fix_stack_size
        ^ param_setup_tp
        ^ Printf.sprintf "\taddi s0, sp, %d\n" fix_stack_size
      in
      (prologue, param_setup)
    else (Printf.sprintf "%s:\n" f.name, "")
  in

  let body_code =
    f.blocks
    |> List.map (fun blk ->
           compile_block blk reg_map need_stack callee_restore_code)
    |> String.concat ""
  in

  let epilogue =
    if need_stack then
      if not (String.ends_with ~suffix:"\tret\n" body_code) then
        callee_restore_code
        ^ Printf.sprintf
            "\tlw ra, %d(s0)\n\tlw s0, %d(sp)\n\taddi sp, sp, %d\n\tret\n"
            (get_stack_offset "ra") (fix_stack_size - 4) fix_stack_size
      else ""
    else if not (String.ends_with ~suffix:"\tret\n" body_code) then "\tret\n"
    else ""
  in
  (* 暂时放外面
  let epilogue = callee_restore_code ^epilogue in *)

  prologue ^ param_setup ^ body_code ^ epilogue

let compile_program (prog : ir_program) (print_alloc : bool) : string =
  let prologue = ".text\n .globl main\n" in
  let body_asm =
    match prog with
    | Ir_funcs funcs -> List.map compile_func funcs |> String.concat "\n"
    | Ir_funcs_o funcs_o ->
        List.map (fun f -> compile_func_o f print_alloc) funcs_o
        |> String.concat "\n"
  in
  prologue ^ body_asm
